/*
* Copyright (C) 2010-2012 Magnus Skjegstad
* 
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU General Public License
* as published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
* 
* The latest version of this library can be downloaded from http://mist-pubsub.googlecode.com
* 
*/
package com.skjegstad.mist2.modules.udp.unicast;

import com.skjegstad.simpleevents.SimpleAsyncEventListener;
import com.skjegstad.simpleevents.SimpleEvent;
import com.skjegstad.simpleevents.SimpleEventListener;
import com.skjegstad.simpleevents.SimpleEventPipe;
import com.skjegstad.mist2.IMistNetworkModule;
import com.skjegstad.mist2.MistUtil;
import com.skjegstad.mist2.NetworkModuleHelper;
import com.skjegstad.mist2.datatypes.BloomfilterAcknowledgment;
import com.skjegstad.mist2.datatypes.ConnectionIdentifier;
import com.skjegstad.mist2.datatypes.DataFragment;
import com.skjegstad.mist2.datatypes.NodeIdentifier;
import com.skjegstad.mist2.datatypes.SimpleEntry;
import com.skjegstad.mist2.datatypes.TopicIdentifier;
import com.skjegstad.mist2.datatypes.TopicMap;
import com.skjegstad.mist2.datatypes.TopicMapEntry;
import com.skjegstad.mist2.events.ExceptionEvent;
import com.skjegstad.mist2.events.NewDataElementReceivedEvent;
import com.skjegstad.mist2.events.NewNodeDiscoveredEvent;
import com.skjegstad.mist2.events.RegisterNodeEvent;
import com.skjegstad.mist2.events.RemoveNodeEvent;
import com.skjegstad.mist2.exceptions.MistException;
import com.skjegstad.mist2.managers.MistEventManager;
import com.skjegstad.mist2.modules.udp.BeaconMessage;
import com.skjegstad.mist2.modules.udp.DataFragmentMessage;
import com.skjegstad.mist2.modules.udp.MessageContainer;
import java.io.Externalizable;
import java.net.DatagramPacket;
import java.net.InetAddress;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Random;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;
import no.ffi.xudp.interfaces.XUdpThreadExceptionCallback;
import no.ffi.xudp.interfaces.XUdpThreadStoppedCallback;
import no.ffi.xudp.threads.XUdpUnicastListener;
import no.ffi.xudp.threads.XUdpUnicastSender;

/**
 *
 * @author Magnus Skjegstad
 */
public class UDPUnicastModule implements IMistNetworkModule, XUdpThreadExceptionCallback, XUdpThreadStoppedCallback {
    private NetworkModuleHelper networkModuleHelper;
    private int fragmentSize;
    private ConnectionIdentifier connectionIdentifier;
    private NodeIdentifier nodeIdentifier;
    private MistEventManager eventManager;
    private SimpleEventPipe localEventPipe;
    private Random r = new Random();
    private ExecutorService executor;
    private BlockingQueue<DatagramPacket> receiveQueue;
    private BlockingQueue<DatagramPacket> sendQueue;
    private BlockingQueue<DatagramPacket> retrySendQueue;
    private long lastBeaconSent = 0;
    private long hashesInLastAck = -1; // update extrabits when the number of hashes is the same in two bloom filters
    private byte extraBits = 17;
    private int k = 1;

    private final int periodicBeaconTime = 60000; // every minute
    private final int timerTick = 250;
    private boolean shouldSendBeacon = true;

    private final int listenPort;

    private final Map<NodeIdentifier, Long> lastSentFragments = new HashMap();
    private final Map<NodeIdentifier, Long> nodeRetries = new HashMap();
    private final ConcurrentMap<NodeIdentifier, Entry<InetAddress, Integer>> connections = new ConcurrentHashMap();

    private int maxNodeRetries = 3;
    private int minRetryDelay = 1000;

    private XUdpUnicastListener listener;
    private XUdpUnicastSender sender;
    private XUdpUnicastSender retrySender;

    public UDPUnicastModule(int listenPort) {
        this.listenPort = listenPort;
        this.fragmentSize = 64000-100;

    }

    protected void setupUDPUnicastThreads() {
        receiveQueue = new LinkedBlockingQueue<DatagramPacket>();
        sendQueue = new LinkedBlockingQueue<DatagramPacket>();
        retrySendQueue = new LinkedBlockingQueue<DatagramPacket>();

        listener = new XUdpUnicastListener(receiveQueue, null, listenPort, null);
        sender = new XUdpUnicastSender(sendQueue, 0, null);
        retrySender = new XUdpUnicastSender(retrySendQueue, 2, null); // 2 retries
    }

    public void tryToConnectToNeighbor(InetAddress host, int port) throws MistException {
        sendBeacon(createBeacon(), host, port, true);
    }

    public void start(final NetworkModuleHelper networkModuleHelper) throws MistException {
        executor = Executors.newCachedThreadPool();
        eventManager = networkModuleHelper.getMistEventManager();
        localEventPipe = eventManager.getUniquePipe("UDP_UNICAST");
        this.networkModuleHelper = networkModuleHelper;

        setupUDPUnicastThreads();        
        
        listener.registerThreadExceptionCallback(this);
        listener.registerThreadStoppedCallback(this);
        
        executor.submit(listener);

        sender.registerExceptionCallback(this);
        sender.registerThreadStoppedCallback(this);
        
        executor.submit(sender);
        
        retrySender.registerExceptionCallback(this);
        retrySender.registerThreadStoppedCallback(this);
        
        executor.submit(retrySender);

        executor.submit(new Runnable() {
            public void run() {
                try {
                    while (!Thread.interrupted()) {
                        DatagramPacket p = receiveQueue.take();
                        receiveUDPPacket(p);
                    }
                } catch (InterruptedException ex) {
                    eventManager.throwException(new MistException("Interrupted while listening for packets", ex));                    
                    return;
                } catch (Exception ex) {
                    eventManager.throwException(new MistException("Unexpected exception while listening for udp packets", ex));
                    return;
                }
            }
        });

        registerEvents();
        
        return;
    }

    public void stop() {
        executor.shutdownNow();
        // TODO should shut own local event listeners
    }

    private int randomMs(int ms, double perm) {
        double part = ms * perm;
        int ret = (int)Math.round(part + (r.nextDouble() * ms));
        //System.err.println(ret);
        return ret;
    }

    private DatagramPacket classToDatagram(Externalizable c, InetAddress destinationAddress, int destinationPort) {
        try {
            return MistUtil.classToDatagram(c, destinationAddress, destinationPort);
        } catch (Exception ex) {
            eventManager.throwException(new MistException("Unable to serialize class", ex));
        }
        return null;
    }
    
    private void sendMessage(Externalizable message, InetAddress dstAddress, int dstPort) throws MistException {
        try {
            sendQueue.add(classToDatagram(new MessageContainer(message), dstAddress, dstPort));
        } catch (ClassNotFoundException ex) {
            throw new MistException("Unable to send message", ex);
        }
    }

    private void sendBeacon(BeaconMessage bm, InetAddress destinationAddress, int destinationPort, boolean retry) throws MistException {
        DatagramPacket p;
        try {
            p = classToDatagram(new MessageContainer(bm), destinationAddress, destinationPort);
        } catch (ClassNotFoundException ex) {
            throw new MistException("Unable to serialize message", ex);
        }
        try {
            networkModuleHelper.debug("Sending beacon to " + p.getAddress() + ":" + p.getPort() + ", retry=" + retry);
            if (retry)
                retrySendQueue.put(p);
            else
                sendQueue.put(p);
            
        } catch (Exception ex) {
            throw new MistException("Got exception while putting beacon in send queue", ex);
        }
    }

    private void sendBeaconToAll() throws MistException {
        synchronized(connections) {
            BeaconMessage bm = createBeacon();
            for (Entry<InetAddress, Integer>e : connections.values())
                sendBeacon(bm, e.getKey(), e.getValue(), false);
        }

    }

    private BeaconMessage createBeacon() throws MistException {
        try {
            updateK();
            networkModuleHelper.debug("Creating beacon w/k=" + k);
            TopicMap topics = networkModuleHelper.getActiveTopics();
            Collection<Long> hashes = networkModuleHelper.getHashesToAck(networkModuleHelper.getActiveTopics().keySet(), false, false); // include expired and invalid items in beacon until they are cleaned up

            if (hashes.size() == hashesInLastAck && !hashes.isEmpty())
                extraBits = (byte) ((extraBits + k) % networkModuleHelper.getMaximumExtraBits());
            else
                hashesInLastAck = hashes.size();

            if (hashes.isEmpty())
                networkModuleHelper.debug("Warning: Sending empty ack");

            BeaconMessage bm;
            int msg_len;
            do {
                bm = new BeaconMessage(networkModuleHelper.getMyNodeIdentifier(),
                    topics, 
                    listenPort,
                    networkModuleHelper.getMaximumAge(),
                    new BloomfilterAcknowledgment(hashes, extraBits, k));
                msg_len = classToDatagram(new MessageContainer(bm), InetAddress.getLocalHost(), 1234).getLength();
                if (msg_len > fragmentSize) {
                    networkModuleHelper.debug("Beacon to large. Reducing k.");
                    k--;
                }
            } while (msg_len > fragmentSize && k > 1);

            lastBeaconSent = System.currentTimeMillis();

            return bm;
        } catch (Exception ex) {
            throw new MistException("Unable to create beacon", ex);
        }
    }

    private void sendDataFragments() throws MistException {            
        try {                        

            // synchronize iterator, just in case
            Collection<NodeIdentifier> neighbors = networkModuleHelper.getNodes();
            synchronized(neighbors) {
                for (NodeIdentifier node : neighbors) {
                    // make sure we know the ip and port of the node
                    Entry<InetAddress, Integer> dest = connections.get(node);
                    if (dest == null) {
                        networkModuleHelper.debug("IP of " + node + " unknown. Skipping.");
                        continue;
                    }
                    
                    // remove node if it has too many retries
                    if (nodeRetries.containsKey(node) && nodeRetries.get(node) >= maxNodeRetries) {
                        networkModuleHelper.debug(("Node " + node + " has not responded for " + maxNodeRetries + ". Removing as neighbor."));
                        networkModuleHelper.removeNeighborNode(node);
                        nodeRetries.remove(node);
                        lastSentFragments.remove(node);
                        connections.remove(node);
                        continue;
                    }
                    // don't send if it is less than minRetryDelay milliseconds since last time for this node
                    if (lastSentFragments.containsKey(node) &&
                            System.currentTimeMillis() - lastSentFragments.get(node) < minRetryDelay) {
                        networkModuleHelper.debug("Waiting to send data to node " + node + ". Ack not received yet.");
                        continue;
                    }

                    // Send the hashes
                    List<Long> hashes = new ArrayList(networkModuleHelper.getHashesMissingInNode(node, true, true )); // exclude invalid and expired from sent messages
                    if (hashes != null && !hashes.isEmpty()) {
                        lastSentFragments.put(node, System.currentTimeMillis());
                        if (nodeRetries.get(node) != null) {
                            nodeRetries.put(node, nodeRetries.get(node)+1);
                        } else
                            nodeRetries.put(node, new Long(1));
                    } else
                        continue;
                    
                    Collections.shuffle(hashes); // shuffle them
                    
                    Collection<DataFragment> fragments = networkModuleHelper.getDataFragments(hashes, true, true); // exclude invalid and expired data elements
                    if (fragments.isEmpty())
                        continue;
                    networkModuleHelper.debug("About to send " + fragments.size() + " data fragments to " + dest.getKey() + ":" + dest.getValue());

                    for (DataFragment f : fragments) {
                        try {
                            sendMessage(new DataFragmentMessage(f), dest.getKey(), dest.getValue());
                        } catch (Exception ex) {
                            throw new MistException("Unable to send data fragment " + f.toString(), ex);
                        }
                    }
                }
            }
                    } catch (Exception ex) {
            throw new MistException("Unable to send data fragments", ex);
        }
    }

    private void receiveUDPPacket(DatagramPacket p) {
        try {                    
            MessageContainer mc = new MessageContainer();
            MistUtil.datagramToClass(p.getData(), mc);

            if (mc.getMessageObject() instanceof BeaconMessage) {
                BeaconMessage bc = (BeaconMessage) mc.getMessageObject();
                receiveBeaconMessage(bc, p.getAddress());
            } else
            if (mc.getMessageObject() instanceof DataFragmentMessage) {
                DataFragmentMessage m = (DataFragmentMessage) mc.getMessageObject();

                // skip our own fragments
                if (m.getFragment().getSourceNode().equals(networkModuleHelper.getMyNodeIdentifier()))
                    return;

                networkModuleHelper.addDataFragments(Collections.singletonList(m.getFragment()));
                shouldSendBeacon = true;
            } else
                eventManager.throwException(new MistException("Unknown message type received!"));

        } catch (Exception ex) {
            eventManager.throwException(new MistException("Unable to deserialize datagram packet", ex));
            return;
        }
    }
    
    private void receiveBeaconMessage(BeaconMessage m, InetAddress sourceAddress) {
        try {
            if (!m.getNodeIdentifier().equals(networkModuleHelper.getMyNodeIdentifier())) {
                String s = "";
                for (int i = 0; i < m.getAcknowledgements().getData().length; i++) {
                    if (!s.isEmpty())
                        s = s + ".";
                    s = s + m.getAcknowledgements().getData()[i];
                }

                // store source ip and listen port
                connections.put(m.getNodeIdentifier(),
                        new SimpleEntry<InetAddress, Integer>(sourceAddress, m.getContactPort()));

                networkModuleHelper.debug("Got beacon from " + m.getNodeIdentifier() + "@" + sourceAddress + ":" + m.getContactPort() +
                        ", len: " + m.getAcknowledgements().getLengthInBits() +
                        " bytes: " + m.getAcknowledgements().getData().length  +
                        ", ackdata: \"" + s + "\"");

                for (Entry<TopicIdentifier, TopicMapEntry> et : m.getSubscribeToTopics().entrySet())
                        networkModuleHelper.debug("Got topic in beacon " + et.getKey().getHashes() + ", ttl:" + et.getValue().getTtl() + ", checksum: " + et.getValue().getChecksum());

                // reset retry counter
                nodeRetries.put(m.getNodeIdentifier(), new Long(0));
                lastSentFragments.remove(m.getNodeIdentifier());

                // register last seen
                eventManager.getNodeManagerEventPipe().trigger(new RegisterNodeEvent(connectionIdentifier,
                    m.getNodeIdentifier(),
                    sourceAddress,
                    m.getSubscribeToTopics(),
                    m.getAcknowledgements()));
            }
        } catch (Exception ex) {
            eventManager.throwException(new MistException("Unexpected exception while receiving beacon", ex));
        }
    }
    
    private void registerEvents() {

        // On new data element registered in data store
        eventManager.getDataManagerEventPipe().listen(new SimpleEventListener() {
            public void receiveEvent(SimpleEvent event) {
                if (event instanceof NewDataElementReceivedEvent) {
                    shouldSendBeacon = true;
                }
            }
        });
        
        // on new node discovered
        eventManager.getNodeManagerEventPipe().listen(new SimpleEventListener() {
            public void receiveEvent(SimpleEvent event) {
                if (event instanceof NewNodeDiscoveredEvent) {
                    shouldSendBeacon = true;
                }
            }
        });

        // On periodic beacon event
        localEventPipe.listen(new SimpleAsyncEventListener() {
            public void receiveEvent(SimpleEvent event) {
                if (event instanceof PeriodicBeaconSendEvent) {
                    // this event should only trigger when there is minimum timerTick ms since we sent something last
                    // since beacons may be sent as a normal reply as well, we have to keep track of them
                    long timeDiff = System.currentTimeMillis() - lastBeaconSent;
                    if (timeDiff < periodicBeaconTime) {
                        try {
                            Thread.sleep(periodicBeaconTime - timeDiff);
                        } catch (InterruptedException ex) {
                            return;
                        }
                        localEventPipe.trigger(event);
                        return;
                    }

                    try {
                        networkModuleHelper.debug(("--- Sending periodic beacon after timer expired --- "));
                        sendBeaconToAll();
                    } catch (MistException ex) {
                        eventManager.throwException(new MistException("Unable to send periodic beacon, retrying in 10 seconds", ex));
                        try {
                            Thread.sleep(10 * 1000);
                        } catch (InterruptedException ex1) {
                            return;
                        }
                        localEventPipe.trigger(event);
                        return;
                    }
                    try {
                        // trigger new beacon in the future
                        Thread.sleep(randomMs(timerTick, 0.5));
                    } catch (InterruptedException ex) {
                        return;
                    }
                    localEventPipe.trigger(event);
                }
            }
        });
        localEventPipe.trigger(new PeriodicBeaconSendEvent());
                      
        // On periodic neighbor update event
        localEventPipe.listen(new SimpleAsyncEventListener() {
            public void receiveEvent(SimpleEvent event) {
                if (event instanceof PeriodicNeighborUpdateEvent) {
                    try {
                        // remove expired nodes
                        Collection<NodeIdentifier> expired = networkModuleHelper.getExpiredNodes(periodicBeaconTime * 3);
                        if (expired != null && !expired.isEmpty()) {
                            eventManager.getNodeManagerEventPipe().trigger(new RemoveNodeEvent(expired, connectionIdentifier));
                            // remove stored IPs and ports
                            for (NodeIdentifier n : expired)
                                connections.remove(n);
                        }

                        int del = sendQueue.size();
                        sendQueue.clear(); // remove unsent stuff

                        if (del > 0)
                            networkModuleHelper.debug("Discarded approx. " + del + " items from sendQueue");

                        if (shouldSendBeacon) {
                            sendBeaconToAll();
                            shouldSendBeacon = false;
                        }

                        sendDataFragments();

                    } catch (Exception ex) {
                        eventManager.throwException(new MistException("Unable to send neighbor update. Update aborted.", ex));
                    }
                    try {
                        Thread.sleep(randomMs(timerTick, 0.5));
                    } catch (InterruptedException ex) {
                        return;
                    }
                    localEventPipe.trigger(event);                    
                }
            }
        });
        localEventPipe.trigger(new PeriodicNeighborUpdateEvent());                
    }
    
    private void updateK() {
        int nodes = networkModuleHelper.getNodes().size(); // +1 is us
        if (nodes > 0) {
            int j = (int) Math.ceil(8.0 / nodes);
            if (j < 1)
                k = 1;
            else
                k = j;
        } else
            k = 8;
    }

    public int getFragmentSize() {
        return fragmentSize;
    }

    public void setConnectionIdentifier(final ConnectionIdentifier connectionIdentifier) {
        this.connectionIdentifier = connectionIdentifier;
    }

    public ConnectionIdentifier getConnectionIdentifier() {
        return this.connectionIdentifier;
    }

    public void setNodeIdentifier(NodeIdentifier nodeIdentifier) {
        this.nodeIdentifier = nodeIdentifier;
    }

    public NodeIdentifier getNodeIdentifier() {
        return nodeIdentifier;
    }

    public void threadExceptionCallback(Runnable runnable, Exception ex) {
        eventManager.throwException(new MistException("Thread exception from runnable " + runnable.toString(), ex));
    }

    public void threadStoppedCallback(Runnable thread) {
        eventManager.throwException(new MistException("Thread stopped: " + thread.toString()));
    }

    public int getListenPort() {
        return listenPort;
    }


    /**
     *
     *
     * EVENTS
     *
     *
     *
     */

       
    public class PeriodicBeaconSendEvent implements SimpleEvent {}

    public class PeriodicNeighborUpdateEvent implements SimpleEvent {}

}
